#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# -*- coding: utf-8 -*-

import os
import gym
import numpy as np
import matplotlib.pyplot as plt
import parl
from parl.utils import logger  # 日志打印工具
import gym_game2048
from scipy import io
from datetime import datetime
from tkinter import *
# import pickle
import time
import threading

from model import Model_2048
from algorithm import DQN_2048
from agent import Agent_2048
from replay_memory import ReplayMemory

GAMMA = 0.995
LEARNING_RATE = 0.001  # 学习率

SIZE = 500
GRID_LEN = 4
GRID_PADDING = 1

BACKGROUND_COLOR_GAME = "#92877d"
BACKGROUND_COLOR_CELL_EMPTY = "#9e948a"
BACKGROUND_COLOR_DICT = {
    2: "#eee4da",
    4: "#ede0c8",
    8: "#f2b179",
    16: "#f59563",
    32: "#f67c5f",
    64: "#f65e3b",
    128: "#edcf72",
    256: "#edcc61",
    512: "#edc850",
    1024: "#edc53f",
    2048: "#edc22e",
}

CELL_COLOR_DICT = {
    2: "#776e65",
    4: "#776e65",
    8: "#f9f6f2",
    16: "#f9f6f2",
    32: "#f9f6f2",
    64: "#f9f6f2",
    128: "#f9f6f2",
    256: "#f9f6f2",
    512: "#f9f6f2",
    1024: "#f9f6f2",
    2048: "#f9f6f2",
}

FONT = ("Fira Code", 20, "bold")


# 多线程
def thread_it(func, *args):
    '''将函数放入线程中执行'''
    # 创建线程
    t = threading.Thread(target=func, args=args)
    # 守护线程
    t.setDaemon(True)
    # 启动线程
    t.start()


class GameGrid(Frame):
    def __init__(self, boards):
        """
        Implementing game grid.

        Show a Tkinter window with the game board.

        Parameters
        ----------
        boards : List
            List of boards of a played game.
        """
        Frame.__init__(self)

        self.grid()
        self.master.title("2048")
        # self.master.maxsize(1000, 1000)
        # self.boards = boards

        # self.matrix = self.boards[0]
        self.matrix = boards[0]
        self.grid_cells = []
        self.init_grid()
        self.update_grid_cells()

        self.wait_visibility()

    def init_grid(self):
        background = Frame(self, bg=BACKGROUND_COLOR_GAME, width=SIZE, height=SIZE)
        background.grid()
        for i in range(GRID_LEN):
            grid_row = []
            for j in range(GRID_LEN):
                cell = Frame(background, bg=BACKGROUND_COLOR_CELL_EMPTY, width=SIZE / GRID_LEN, height=SIZE / GRID_LEN)
                cell.grid(row=i, column=j, padx=GRID_PADDING, pady=GRID_PADDING)
                t = Label(
                    master=cell, text="", bg=BACKGROUND_COLOR_CELL_EMPTY, justify=CENTER, font=FONT, width=4, height=2
                )
                t.grid()
                grid_row.append(t)

            self.grid_cells.append(grid_row)

    def update_grid_cells(self):
        for i in range(GRID_LEN):
            for j in range(GRID_LEN):
                new_number = self.matrix[i][j]
                if new_number == 0:
                    self.grid_cells[i][j].configure(text="", bg=BACKGROUND_COLOR_CELL_EMPTY)
                else:
                    self.grid_cells[i][j].configure(
                        text=str(new_number), bg=BACKGROUND_COLOR_DICT[new_number], fg=CELL_COLOR_DICT[new_number]
                    )
        self.update_idletasks()

    def make_move(self, boards):

        # for board in self.boards:
        for board in boards:
            time.sleep(0.05)
            self.matrix = board
            self.update_grid_cells()
        print('end')
        time.sleep(1)
        sys.exit()


def remove_repeated_boards(boards):
    "Remove remove repeat boards from list. Some movements are invalid, so they don't change the board."
    no_repeated_boards = []
    last_board = None
    for i, board in enumerate(boards):

        if isinstance(board, np.ndarray):
            if np.array_equal(last_board, board) is True:
                pass
            else:
                no_repeated_boards.append(board)
                last_board = board

        else:
            last_board = board
            no_repeated_boards.append(board)

    return no_repeated_boards


# 评估 agent
def evaluate(env, agent, render=False):
    obs = env.reset()
    episode_reward = 0
    total_step = 0
    board = []
    action = []
    score = []
    while True:
        act = agent.predict(np.transpose(obs, (2, 0, 1)))  # 预测动作，只选最优动作
        old_board = env.get_board()
        obs, reward, done, _ = env.step(act)
        new_board = env.get_board()
        flag = 0
        while np.array_equal(new_board, old_board):
            if flag == 0:
                Q = agent.predict_prob(np.transpose(obs, (2, 0, 1)))
            # trick：如果board不变，选取Q次大的动作
            Q[act] = 0
            act = np.argmax(Q)
            obs, reward, done, _ = env.step(act)
            old_board = new_board
            new_board = env.get_board()
            flag += 1
            if flag > 3:
                break  # done
        episode_reward += reward
        total_step += 1
        board.append(new_board)
        action.append(act)
        score.append(reward)
        if render:
            print(total_step)
            print(act)
            print(reward)
            print(episode_reward)
            print(new_board)
            env.render()
        if done:
            break
    return board, action, score, env.get_board().max()


def main():
    print('hello 2048!')
    env = gym.make("game2048-v0", board_size=4, seed=None, binary=True, extractor="cnn", penalty=0)
    act_dim = env.action_space.n
    obs_shape = list(env.observation_space.shape)

    # 根据parl框架构建agent
    model = Model_2048(act_dim=act_dim)
    algorithm = DQN_2048(model, act_dim=act_dim, gamma=GAMMA, lr=LEARNING_RATE)
    agent = Agent_2048(
        algorithm,
        obs_dim=obs_shape,
        act_dim=act_dim,
        e_greed=0.1,  # 有一定概率随机选取动作，探索
        e_greed_decrement=1e-6)  # 随着训练逐步收敛，探索的程度慢慢降低

    # 加载模型
    # agent.restore('model/steps_13860_9479.6_19800.ckpt')
    agent.restore('model/steps_14228_10063.6_1200.ckpt')


    max_episode = 100
    episode = 0
    best_score = 0
    res_max = np.zeros(8)
    testReward = []
    while episode < max_episode:
        episode += 1
        # train part
        board, action, score, maxnum = evaluate(env, agent, render=False)
        res_score = sum(score)
        step = len(score)
        testReward.append([episode, res_score, step])

        if maxnum < 64:
            res_max[0] += 1
        elif maxnum == 64:
            res_max[1] += 1
        elif maxnum == 128:
            res_max[2] += 1
        elif maxnum == 256:
            res_max[3] += 1
        elif maxnum == 512:
            res_max[4] += 1
        elif maxnum == 1024:
            res_max[5] += 1
        elif maxnum == 2048:
            res_max[6] += 1
        else:
            res_max[7] += 1

        if res_score > best_score:
            best_score = res_score
            best_board = board
            best_action = action
        logger.info('episode:{}    Best score:{}   MaxNum:{}   Step:{}'.format(episode, res_score, maxnum, step))
    print(res_max / max_episode)
    print(best_score)
    return best_board, best_action


def show(best_board):
    show_boards = remove_repeated_boards(best_board)
    root = Tk()
    gamegrid = GameGrid(show_boards)
    # 使用多线程显示，否则会‘未响应’
    thread_it(gamegrid.make_move, show_boards)
    root.mainloop()
    root.destory()


if __name__ == '__main__':
    best_board, best_action = main()
    np.save("best_game.npy", best_board)
    bs = np.load("best_game.npy")
    show(bs)
